#include <bsl_params.h>
#include <bsl_errno.h>
#include <bsl_err_internal.h>
#include <crypt_algid.h>
#include <crypt_errno.h>
#include <securec.h>

#include <wd.h>
#include <wd_digest.h>
#include <wd_bmm.h>

#include "digests.h"
#include "err.h"

#include <stdint.h>
#include <stdio.h>

void KAE_DGST_Free(KAE_DGST_Ctx* ctx)
{
    if (ctx == NULL) {
        return ;
    }

    if (ctx->pool != NULL && ctx->opData.in != NULL) {
        wd_free_blk(ctx->pool, ctx->opData.in);
    }

    if (ctx->pool != NULL && ctx->opData.out != NULL) {
        wd_free_blk(ctx->pool, ctx->opData.out);
    }

    if (ctx->pool != NULL) {
        wd_blkpool_destroy(ctx->pool);
    }

    wd_release_queue(&(ctx->queue));

    if (ctx->digestCtx != NULL) {
        wcrypto_del_digest_ctx(ctx->digestCtx);
    }

    free(ctx);
}

KAE_DGST_Ctx* KAE_DGST_NewCtx(CRYPT_MD_AlgId md_alg)
{
    KAE_DGST_Ctx* ctx = (KAE_DGST_Ctx*) malloc(sizeof(KAE_DGST_Ctx));
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }

    (void) memset_s(ctx, sizeof(KAE_DGST_Ctx), 0, sizeof(KAE_DGST_Ctx));
    ctx->queue.capa.alg = "digest";
    int32_t ret = wd_request_queue(&(ctx->queue));
    if (ret != KAE_WD_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_REQUEST_FAIL);
        return NULL;
    }

    struct wd_blkpool_setup pool_setup;
    (void) memset_s(&pool_setup, sizeof(pool_setup), 0, sizeof(pool_setup));
    pool_setup.block_size = DGST_BLOCK_SIZE;

    switch (md_alg)
    {
        case CRYPT_MD_SM3:
        {
            pool_setup.block_num = SM3_BLOCK_NUM;
            ctx->digestSize = SM3_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SM3;
            break;
        }
        case CRYPT_MD_MD5:
        {
            pool_setup.block_num = MD5_BLOCK_NUM;
            ctx->digestSize = MD5_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_MD5;
            break;
        }
        case CRYPT_MD_SHA1:
        {
            pool_setup.block_num = SHA1_BLOCK_NUM;
            ctx->digestSize = SHA1_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA1;
            break;
        }
        case CRYPT_MD_SHA256:
        {
            pool_setup.block_num = SHA256_BLOCK_NUM;
            ctx->digestSize = SHA256_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA256;
            break;
        }
        case CRYPT_MD_SHA224:
        {
            pool_setup.block_num = SHA224_BLOCK_NUM;
            ctx->digestSize = SHA224_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA224;
            break;
        }
        case CRYPT_MD_SHA384:
        {
            pool_setup.block_num = SHA384_BLOCK_NUM;
            ctx->digestSize = SHA384_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA384;
            break;
        }
        case CRYPT_MD_SHA512:
        {
            pool_setup.block_num = SHA512_BLOCK_NUM;
            ctx->digestSize = SHA512_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA512;
            break;
        }
        default:
        {
            BSL_ERR_PUSH_ERROR(KAE_WD_NOT_SUPPORT);
            goto ERR;
        }
    }
    ctx->id = md_alg;

    pool_setup.align_size = 64;
    ctx->pool = wd_blkpool_create(&ctx->queue, &pool_setup);
    if (ctx->pool == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }

    ctx->setup.mode = WCRYPTO_DIGEST_NORMAL;
    ctx->setup.br.alloc = (void*) wd_alloc_blk;
    ctx->setup.br.free = (void*) wd_free_blk;
    ctx->setup.br.iova_map = (void*) wd_blk_iova_map;
    ctx->setup.br.iova_unmap = (void*) wd_blk_iova_unmap;
    ctx->setup.br.usr = ctx->pool;
    ctx->digestCtx = wcrypto_create_digest_ctx(&ctx->queue, &ctx->setup);
    if (ctx->digestCtx == NULL) {
        BSL_ERR_PUSH_ERROR(KAE_WD_CREATE_CTX_FAIL);
        goto ERR;
    }

    ctx->opData.in = wd_alloc_blk(ctx->pool);
    if (ctx->opData.in == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }

    ctx->opData.out = wd_alloc_blk(ctx->pool);
    if (ctx->opData.out == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }

    ctx->opData.out_bytes = ctx->digestSize;
    (void) memset_s(ctx->opData.out, ctx->digestSize, 0, ctx->digestSize);
    ctx->opData.has_next = false;

    return ctx;
ERR:
    KAE_DGST_Free(ctx);
    return NULL;
}

int32_t KAE_DGST_Init(KAE_DGST_Ctx* ctx, BSL_Param* param)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    (void) param;

    return CRYPT_SUCCESS;
}

int32_t KAE_DGST_Update(KAE_DGST_Ctx* ctx, const uint8_t* in, uint32_t len)
{
    if ((ctx == NULL) || (ctx->digestCtx == NULL) || (in == NULL && len != 0)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (len == 0) {
        return CRYPT_SUCCESS;
    }

    int cnt = 0;

    if (ctx->remain > 0)
    {
        if (len + ctx->remain > DGST_CACHE_SIZE)
        {
            (void) memcpy_s(ctx->block + ctx->remain, DGST_CACHE_SIZE - ctx->remain, in, DGST_CACHE_SIZE - ctx->remain);
            in += DGST_CACHE_SIZE - ctx->remain;
            len -= (DGST_CACHE_SIZE - ctx->remain);
            ctx->remain = 0;
            ctx->opData.in_bytes = DGST_CACHE_SIZE;
            ctx->opData.has_next = true;
            (void) memcpy_s(ctx->opData.in, DGST_CACHE_SIZE, ctx->block, ctx->opData.in_bytes);
            if (wcrypto_do_digest(ctx->digestCtx, &ctx->opData, NULL) != CRYPT_SUCCESS) {
                BSL_ERR_PUSH_ERROR(KAE_WD_DO_DGST_FAIL);
                return KAE_WD_DO_DGST_FAIL;
            }
        }
        else
        {
            (void) memcpy_s(ctx->block + ctx->remain, DGST_CACHE_SIZE - ctx->remain, in, len);
            ctx->remain += len;
            return CRYPT_SUCCESS;
        }
    }

    while (len - cnt > DGST_CACHE_SIZE)
    {
        ctx->opData.in_bytes = DGST_CACHE_SIZE;
        ctx->opData.has_next = true;
        (void) memcpy_s(ctx->opData.in, DGST_CACHE_SIZE, in, ctx->opData.in_bytes);
        if (wcrypto_do_digest(ctx->digestCtx, &ctx->opData, NULL) != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(KAE_WD_DO_DGST_FAIL);
            return KAE_WD_DO_DGST_FAIL;
        }
        in += DGST_CACHE_SIZE;
        cnt += DGST_CACHE_SIZE;
    }

    ctx->remain = len - cnt;
    (void) memcpy_s(ctx->block, DGST_CACHE_SIZE, in, ctx->remain);

    return CRYPT_SUCCESS;
}

int32_t KAE_DGST_Final(KAE_DGST_Ctx* ctx, uint8_t* out, uint32_t* len)
{
    if (ctx == NULL || ctx->digestCtx == NULL || (out == NULL) || (len == NULL)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (*len < ctx->digestSize) {
        BSL_ERR_PUSH_ERROR(CRYPT_MD5_OUT_BUFF_LEN_NOT_ENOUGH);
        return CRYPT_MD5_OUT_BUFF_LEN_NOT_ENOUGH;
    }

    ctx->opData.in_bytes = ctx->remain;
    ctx->opData.has_next = false;
    (void) memcpy_s(ctx->opData.in, DGST_CACHE_SIZE, ctx->block, ctx->opData.in_bytes);
    if (wcrypto_do_digest(ctx->digestCtx, &ctx->opData, NULL) != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_DO_DGST_FAIL);
        return KAE_WD_DO_DGST_FAIL;
    }
    (void) memcpy_s(out, ctx->digestSize, ctx->opData.out, ctx->digestSize);
    *len = ctx->digestSize;

    return CRYPT_SUCCESS;
}

void KAE_DGST_Deinit(KAE_DGST_Ctx* ctx)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return ;
    }

    wd_free_blk(ctx->pool, ctx->opData.in);
    wd_free_blk(ctx->pool, ctx->opData.out);

    ctx->opData.in = wd_alloc_blk(ctx->pool);
    if (ctx->opData.in == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return ;
    }

    ctx->opData.out = wd_alloc_blk(ctx->pool);
    if (ctx->opData.out == NULL) {
        wd_free_blk(ctx->pool, ctx->opData.in);
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return ;
    }

    ctx->opData.out_bytes = ctx->digestSize;
    (void) memset_s(ctx->opData.out, ctx->digestSize, 0, ctx->digestSize);
    ctx->opData.has_next = false;

    (void) memset_s(ctx->block, DGST_CACHE_SIZE, 0, DGST_CACHE_SIZE);
    ctx->remain = 0;
}

KAE_DGST_Ctx* KAE_DGST_Dup(const KAE_DGST_Ctx *src)
{
    if (src == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return NULL;
    }

    KAE_DGST_Ctx* newCtx = KAE_DGST_NewCtx(src->id);
    if (newCtx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    return newCtx;
}